package pers.qianyu.module.core.util;

import pers.qianyu.module.core.domain.image.model.TrieNode;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 字典树工具类
 *
 * @author mizzle rain
 * @date 2021-03-17 11:04
 */
public class TrieTreeUtil {
    private final static TrieNode root = new TrieNode();

    public static void insert(String word) {
        int n = word.length();
        char[] arr = word.toCharArray();
        TrieNode p = root;
        for (int i = 0; i < n; i++) {
            if (!p.containsKey(arr[i])) {
                p.put(arr[i], new TrieNode());
            }
            p = p.get(arr[i]);
        }
        p.setEnd(true);
    }

    public static boolean search(String word) {
        TrieNode node = searchPrefix(word);
        return node != null && node.isEnd();
    }

    public static Set<String> searchByPrefix(String prefix) {
        TrieNode node = searchPrefix(prefix);
        if (Objects.isNull(node)) {
            return new HashSet<>();
        }
        Set<String> set = new BackTrack().solve(node);
        return set.stream().map(s -> prefix + s).collect(Collectors.toSet());
    }

    private static class BackTrack {
        HashSet<String> res;

        public Set<String> solve(TrieNode node) {
            StringBuffer path = new StringBuffer();
            res = new HashSet<>();
            dfs(path, node);
            return res;
        }

        private void dfs(StringBuffer path, TrieNode node) {
            if (Objects.isNull(node)) {
                return;
            }
            if (node.isEnd()) {
                res.add(path.toString());
            }
            for (Map.Entry<Character, TrieNode> entry : node.getLinksMap().entrySet()) {
                Character k = entry.getKey();
                TrieNode v = entry.getValue();
                path.append(k);
                dfs(path, v);
                int len = path.length();
                path.delete(len - 1, len);
            }
        }
    }

    public static boolean startsWith(String prefix) {
        TrieNode node = searchPrefix(prefix);
        return node != null;
    }

    private static TrieNode searchPrefix(String word) {
        TrieNode p = root;
        int n = word.length();
        char[] arr = word.toCharArray();
        for (int i = 0; i < n; i++) {
            if (!p.containsKey(arr[i])) {
                return null;
            }
            p = p.get(arr[i]);
        }
        return p;
    }
}
